#version 450
#extension GL_ARB_separate_shader_objects : enable
#extension GL_EXT_nonuniform_qualifier : enable
#extension GL_EXT_scalar_block_layout : enable
#extension GL_GOOGLE_include_directive : enable
#extension GL_EXT_shader_explicit_arithmetic_types_int64 : require
#extension GL_EXT_buffer_reference2 : require
#extension GL_ARB_shader_image_size : require

#include "Common.glsl"

layout (location = 0) in vec2 inUV;

layout (binding = 0) uniform Ubo {
    mat4 projview;
	mat4 viewInverse;
	mat4 projInverse;
    uint frame;
} ubo;

layout (binding = 1, rgba32f) uniform image2D resultImage;
layout (binding = 2, scalar) buffer AllMaterial { Material m[]; } materials;
layout (binding = 3, scalar) buffer AllVertex { Vertex v[]; } vertexs;
layout (binding = 4, scalar) buffer AllTriangle { Triangle t[]; } triangles;
layout (binding = 5, scalar) buffer AllBVHNode { BVHNode n[]; } bvhNodes;
layout (binding = 6) uniform samplerCube samplerCubeMap;


vec3 toneMap = vec3(1.0f / Uncharted2Tonemap(vec3(W)));
vec3 gammaPow = vec3(1.0f / 2.2);

vec3 CubamapSample(vec3 direction)
{
    vec3 color = Uncharted2Tonemap(texture(samplerCubeMap, direction).rgb * 4.5);
    color = color * toneMap;
	color = pow(color, gammaPow);
    return color;
}

HitResult HitTriangle(Ray ray, uint t)
{
    HitResult result;
    result.isHit = false;
    result.distance = FLT_MAX;
    
    uvec3 index = triangles.t[t].index;
    vec3 p1 = vertexs.v[index.x].position;
    vec3 p2 = vertexs.v[index.y].position;
    vec3 p3 = vertexs.v[index.z].position;

    // 这里的法线是平面法线, 不是真正的模型法线
    result.normal = normalize(cross(p2 - p1, p3 - p1));

    // 是否在背面, 在背面的话应该就是在模型内部了, 法线要反转
    if(dot(result.normal, ray.direction) > 0.0)
    {
        result.normal = -result.normal;
    }

    // 和法线垂直, 和平面平行直接返回
    if(abs(dot(result.normal, ray.direction)) < 0.00001)
        return result;

    // 求射线和平面的交点距离t
    result.distance = (dot(result.normal, p1) - dot(ray.origin, result.normal)) / dot(ray.direction, result.normal);
    if(result.distance < 0.0005)
        return result;

    result.point = ray.origin + ray.direction * result.distance;

    vec3 c1 = cross(p2 - p1, result.point - p1);
    vec3 c2 = cross(p3 - p2, result.point - p2);
    vec3 c3 = cross(p1 - p3, result.point - p3);
    bool r1 = (dot(c1, result.normal) > 0 && dot(c2, result.normal) > 0 && dot(c3, result.normal) > 0);
    bool r2 = (dot(c1, result.normal) < 0 && dot(c2, result.normal) < 0 && dot(c3, result.normal) < 0);

    if(r1 || r2)
    {
        result.isHit = true;
        result.view = ray.direction;
    }

    return result;
}

HitResult HitTriangles(Ray ray, uint l, uint r)
{
    HitResult result;
    result.isHit = false;
    result.distance = FLT_MAX;
    uint t;

    for(uint i = l; i <= r; ++i)
    {
        HitResult res = HitTriangle(ray, i);
        if(res.isHit && res.distance < result.distance)
        {
            result = res;
            t = i;
        }
    }

    result.triangle = t;
    
    return result;
}

float HitAABB(Ray ray, vec3 aa, vec3 bb)
{
    vec3 inPoint = (aa - ray.origin) / ray.direction;
	vec3 outPoint = (bb - ray.origin) / ray.direction;

	vec3 tmax = max(inPoint, outPoint);
	vec3 tmin = min(inPoint, outPoint);

	float t0 = max(tmin.x, max(tmin.y, tmin.z));
	float t1 = min(tmax.x, min(tmax.y, tmax.z));

	return (t1 >= t0) ? (t0 > 0.f ? t0 : t1) : -1;
}

HitResult HitBVH(Ray ray)
{
    HitResult result;
    result.isHit = false;
    result.distance = FLT_MAX;


    uint stack[256];
    uint sp = 0;
    stack[sp++] = 0;


    while(sp > 0)
    {
        BVHNode node = bvhNodes.n[stack[--sp]];

        if(node.n > 0)
        {
            HitResult res = HitTriangles(ray, node.l, node.l + node.n - 1);
            if(res.isHit && res.distance < result.distance)
                result = res;
            continue;
        }

        float dL = FLT_MAX;
        float dR = FLT_MAX;

        if(node.left > 0)
        {
            BVHNode leftNode = bvhNodes.n[node.left];
            dL = HitAABB(ray, leftNode.aa, leftNode.bb);
        }

        if(node.right > 0)
        {
            BVHNode rightNode = bvhNodes.n[node.right];
            dR = HitAABB(ray, rightNode.aa, rightNode.bb);
        }

        if(dL > 0.0 && dR > 0.0)
        {
            if(dL < dR)
            {
                stack[sp++] = node.left;
                stack[sp++] = node.right;
            }
            else
            {
                stack[sp++] = node.right;
                stack[sp++] = node.left;
            }
        }
        else if(dL > 0.0)
        {
            stack[sp++] = node.left;
        }
        else if(dR > 0.0)
        {
            stack[sp++] = node.right;
        }
    }

    if(result.isHit)
    {
        uvec3 index = triangles.t[result.triangle].index;
        Vertex v1 = vertexs.v[index.x];
        Vertex v2 = vertexs.v[index.y];
        Vertex v3 = vertexs.v[index.z];
        
        // 插值计算空间法线切线和uv
        float alpha = (-(result.point.x-v2.position.x)*(v3.position.y-v2.position.y) + (result.point.y-v2.position.y)*(v3.position.x-v2.position.x)) / (-(v1.position.x-v2.position.x-0.00005)*(v3.position.y-v2.position.y+0.00005) + (v1.position.y-v2.position.y+0.00005)*(v3.position.x-v2.position.x+0.00005));
        float beta  = (-(result.point.x-v3.position.x)*(v1.position.y-v3.position.y) + (result.point.y-v3.position.y)*(v1.position.x-v3.position.x)) / (-(v2.position.x-v3.position.x-0.00005)*(v1.position.y-v3.position.y+0.00005) + (v2.position.y-v3.position.y+0.00005)*(v1.position.x-v3.position.x+0.00005));
        float gama  = 1.0 - alpha - beta;


        // result.normal = normalize(alpha * v1.normal + beta * v2.normal + gama * v3.normal);
        // result.tangent = normalize(alpha * v1.tangent + beta * v2.tangent + gama * v3.tangent);
        // result.bitangent = normalize(alpha * v1.bitangent + beta * v2.bitangent + gama * v3.bitangent);

        result.normal = alpha * v1.normal + beta * v2.normal + gama * v3.normal;
        result.tangent = alpha * v1.tangent + beta * v2.tangent + gama * v3.tangent;
        result.tangent = normalize(result.tangent - dot(result.tangent, result.normal) * result.normal);
        result.bitangent = cross(result.normal, result.tangent);

        result.uv = alpha * v1.uv + beta * v2.uv + gama * v3.uv;

        result.material = triangles.t[result.triangle].material;
    }

    return result;
}

vec3 TracingPath(inout uint seed, HitResult hit, uint depth)
{
    vec3 color = vec3(0.0);
    vec3 history = vec3(1.0);

    for(uint i = 0; i < depth; ++i)
    {
        Ray ray;
        ray.origin = hit.point;
        ray.direction = TurnToTangentSpace(HelfSphereSample(seed), hit.normal, hit.tangent, hit.bitangent);
        // ray.direction = TurnToTangentSpace(ImportanceSample(seed), hit.normal, hit.tangent, hit.bitangent);

        HitResult result = HitBVH(ray);

        float cosTheta = max(dot(ray.direction, hit.normal), 0.0);
        vec3 diffuse = materials.m[hit.material].baseColor / M_PI;

        if(!result.isHit)
        {
            vec3 skyColor = CubamapSample(ray.direction);
            color += history * skyColor * diffuse * cosTheta / HelfSpherePdf;
            break;
        }

        vec3 specular = materials.m[result.material].emissive;
        color += history * specular * diffuse * cosTheta / HelfSpherePdf;

        hit = result;
        history *= diffuse * cosTheta / HelfSpherePdf;
    }

    return color;
}

void main()
{
    ivec2 texUV = imageSize(resultImage);
    texUV.x = int(inUV.x * texUV.x);
    texUV.y = int(inUV.y * texUV.y);
    uint seed = Tea(texUV.x * texUV.y + texUV.y, ubo.frame);

    vec2 uv = inUV * 2.0 - 1.0;
    Ray ray;
	ray.origin = vec3(ubo.viewInverse * vec4(0, 0, 0, 1));
    vec3 target = vec3(ubo.projInverse * vec4(uv.x, uv.y, 1, 1));
    ray.direction = vec3(ubo.viewInverse * vec4(normalize(target), 0));


    vec3 color = vec3(0.0);
    HitResult result = HitBVH(ray);
    if(result.isHit)
    {
        color = materials.m[result.material].emissive;
        vec3 specular = TracingPath(seed, result, 2);
        color += specular;
    }
    else
    {
        // Tone mapping
        color = CubamapSample(ray.direction);
    }

#if 1
    if(ubo.frame > 0)
    {
        float a = 1.0f / float(ubo.frame + 1);
        vec3  oldColor = imageLoad(resultImage, texUV).xyz;
        imageStore(resultImage, texUV, vec4(mix(oldColor, color, a), 1.f));
    }
    else
    {
        imageStore(resultImage, texUV, vec4(color, 1.f));
    }
#else
    imageStore(resultImage, texUV, vec4(color, 1.f));
#endif
}